import tushare as ts
import pandas as pd
import os
from datetime import datetime, timedelta
import time

# 设置token
TOKEN = '58832278507aea5a5a7afa384ef20192a3f751f8a6dfe83e094d60f3'  # 请替换为您的token
ts.set_token(TOKEN)
pro = ts.pro_api()

# 定义数据存储路径
BASE_DIR = 'full_stock_data'
DAILY_DIR = os.path.join(BASE_DIR, 'daily')
BASIC_DIR = os.path.join(BASE_DIR, 'basic')
INDEX_DIR = os.path.join(BASE_DIR, 'index')

def ensure_dirs():
    """确保数据目录存在"""
    for dir_path in [BASE_DIR, DAILY_DIR, BASIC_DIR, INDEX_DIR]:
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)

def get_stock_list():
    """获取所有A股上市公司列表"""
    try:
        # 获取所有交易所的股票
        stocks_list = []
        
        # 上交所（主板 + 科创板）
        stocks_sh = pro.stock_basic(
            exchange='SSE',
            list_status='L',
            fields='ts_code,symbol,name,area,industry,list_date,market'
        )
        stocks_list.append(stocks_sh)
        
        # 深交所（主板 + 创业板）
        stocks_sz = pro.stock_basic(
            exchange='SZSE',
            list_status='L',
            fields='ts_code,symbol,name,area,industry,list_date,market'
        )
        stocks_list.append(stocks_sz)
        
        # 北交所
        stocks_bj = pro.stock_basic(
            exchange='BSE',
            list_status='L',
            fields='ts_code,symbol,name,area,industry,list_date,market'
        )
        stocks_list.append(stocks_bj)
        
        # 合并所有股票列表
        stocks = pd.concat(stocks_list, ignore_index=True)
        
        # 保存股票列表
        stocks.to_csv(os.path.join(BASIC_DIR, 'stock_list.csv'), index=False)
        print(f"成功获取 {len(stocks)} 只股票的基本信息")
        return stocks
    except Exception as e:
        print(f"获取股票列表失败: {str(e)}")
        return None

def download_daily_data(ts_code, start_date='19900101'):
    """下载单个股票的历史数据"""
    try:
        end_date = datetime.now().strftime('%Y%m%d')
        
        # 下载日线数据
        df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        
        # 下载复权因子
        factors = pro.adj_factor(ts_code=ts_code, start_date=start_date, end_date=end_date)
        
        # 合并数据
        if not df.empty and not factors.empty:
            df = pd.merge(df, factors, on=['ts_code', 'trade_date'])
            
            # 计算复权价格
            for price_col in ['open', 'high', 'low', 'close']:
                df[price_col] = df[price_col] * df['adj_factor']
            
            # 按日期排序
            df = df.sort_values('trade_date', ascending=True)
            
            # 计算均线
            for ma in [5, 10, 20, 30, 60, 120, 240, 250]:
                df[f'ma{ma}'] = df['close'].rolling(window=ma).mean()
            
            # 计算成交量均线
            for ma in [5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 120, 150, 200, 250]:
                df[f'vol_ma{ma}'] = df['vol'].rolling(window=ma).mean()
            
            # 保存数据
            file_path = os.path.join(DAILY_DIR, f'{ts_code}.csv')
            df.to_csv(file_path, index=False)
            return True
        return False
    except Exception as e:
        print(f"下载 {ts_code} 的数据失败: {str(e)}")
        return False

def download_basic_data(ts_code, start_date='19900101'):
    """下载单个股票的基本面数据"""
    try:
        end_date = datetime.now().strftime('%Y%m%d')
        
        # 下载日线基本面指标
        df = pro.daily_basic(
            ts_code=ts_code,
            start_date=start_date,
            end_date=end_date,
            fields=[
                'ts_code', 'trade_date', 'close', 'turnover_rate', 
                'turnover_rate_f', 'volume_ratio', 'pe', 'pe_ttm',
                'pb', 'ps', 'ps_ttm', 'dv_ratio', 'dv_ttm',
                'total_share', 'float_share', 'free_share',
                'total_mv', 'circ_mv'
            ]
        )
        
        if not df.empty:
            # 按日期排序
            df = df.sort_values('trade_date', ascending=True)
            
            # 保存数据
            file_path = os.path.join(BASIC_DIR, f'{ts_code}_basic.csv')
            df.to_csv(file_path, index=False)
            return True
        return False
    except Exception as e:
        print(f"下载 {ts_code} 的基本面数据失败: {str(e)}")
        return False

def download_all_data():
    """下载所有股票的历史数据和基本面数据"""
    ensure_dirs()
    
    # 获取股票列表
    stocks = get_stock_list()
    if stocks is None:
        return
    
    total_stocks = len(stocks)
    success_daily = 0
    success_basic = 0
    fail_daily = 0
    fail_basic = 0
    
    print(f"开始下载 {total_stocks} 只股票的数据...")
    
    for idx, row in stocks.iterrows():
        ts_code = row['ts_code']
        print(f"\n处理 {ts_code} ({idx+1}/{total_stocks})...")
        print(f"股票名称: {row['name']}, 市场: {row['market']}")
        
        # 下载日线数据
        print("下载行情数据...")
        if download_daily_data(ts_code):
            success_daily += 1
        else:
            fail_daily += 1
        
        # 下载基本面数据
        print("下载基本面数据...")
        if download_basic_data(ts_code):
            success_basic += 1
        else:
            fail_basic += 1
        
        # 避免频繁请求
        time.sleep(1)
    
    print(f"\n下载完成:")
    print(f"行情数据 - 成功: {success_daily}, 失败: {fail_daily}")
    print(f"基本面数据 - 成功: {success_basic}, 失败: {fail_basic}")

def update_all_data():
    """更新所有股票的数据"""
    if not os.path.exists(BASIC_DIR):
        print("未找到基础数据，请先运行完整下载")
        return
    
    stock_list_file = os.path.join(BASIC_DIR, 'stock_list.csv')
    if not os.path.exists(stock_list_file):
        print("未找到股票列表，请先运行完整下载")
        return
    
    stocks = pd.read_csv(stock_list_file)
    total_stocks = len(stocks)
    
    print(f"开始更新 {total_stocks} 只股票的数据...")
    
    for idx, row in stocks.iterrows():
        ts_code = row['ts_code']
        print(f"\n更新 {ts_code} ({idx+1}/{total_stocks})...")
        
        # 更新日线数据
        daily_file = os.path.join(DAILY_DIR, f'{ts_code}.csv')
        if os.path.exists(daily_file):
            df = pd.read_csv(daily_file)
            if not df.empty:
                last_date = df['trade_date'].max()
                start_date = (pd.to_datetime(str(last_date)) + timedelta(days=1)).strftime('%Y%m%d')
                download_daily_data(ts_code, start_date)
        
        # 更新基本面数据
        basic_file = os.path.join(BASIC_DIR, f'{ts_code}_basic.csv')
        if os.path.exists(basic_file):
            df = pd.read_csv(basic_file)
            if not df.empty:
                last_date = df['trade_date'].max()
                start_date = (pd.to_datetime(str(last_date)) + timedelta(days=1)).strftime('%Y%m%d')
                download_basic_data(ts_code, start_date)
        
        # 避免频繁请求
        time.sleep(1)
    
    print("\n更新完成")

def download_index_data():
    """下载主要指数的历史数据"""
    # 主要指数列表
    index_list = [
        {'ts_code': '000001.SH', 'name': '上证指数'},
        {'ts_code': '399001.SZ', 'name': '深证成指'},
        {'ts_code': '399006.SZ', 'name': '创业板指'},
        {'ts_code': '000001.BJ', 'name': '北证50'},
        {'ts_code': '000016.SH', 'name': '上证50'},
        {'ts_code': '000300.SH', 'name': '沪深300'},
        {'ts_code': '000905.SH', 'name': '中证500'},
        {'ts_code': '000852.SH', 'name': '中证1000'}
    ]
    
    start_date = '19900101'  # 起始日期
    end_date = datetime.now().strftime('%Y%m%d')
    
    print("开始下载指数数据...")
    
    for index in index_list:
        ts_code = index['ts_code']
        name = index['name']
        print(f"\n下载 {name}({ts_code}) 的历史数据...")
        
        try:
            # 下载指数日线数据
            df = pro.index_daily(
                ts_code=ts_code,
                start_date=start_date,
                end_date=end_date,
                fields=[
                    'ts_code', 'trade_date', 'open', 'high', 'low', 
                    'close', 'pre_close', 'change', 'pct_chg', 'vol', 
                    'amount'
                ]
            )
            
            if not df.empty:
                # 按日期排序
                df = df.sort_values('trade_date', ascending=True)
                
                # 计算技术指标
                # 计算均线
                for ma in [5, 10, 20, 30, 60, 120, 240,250]:
                    df[f'ma{ma}'] = df['close'].rolling(window=ma).mean()
                
                # 计算成交量均线
                for vol_ma in [5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 120, 150, 200, 250]:
                    df[f'vol_ma{vol_ma}'] = df['vol'].rolling(window=vol_ma).mean()
                
                # 计算MACD
                exp1 = df['close'].ewm(span=12, adjust=False).mean()
                exp2 = df['close'].ewm(span=26, adjust=False).mean()
                df['macd'] = exp1 - exp2
                df['macd_signal'] = df['macd'].ewm(span=9, adjust=False).mean()
                df['macd_hist'] = df['macd'] - df['macd_signal']
                
                # 计算RSI
                delta = df['close'].diff()
                gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
                loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
                rs = gain / loss
                df['rsi'] = 100 - (100 / (1 + rs))
                
                # 计算布林带
                df['boll_mid'] = df['close'].rolling(window=20).mean()
                df['boll_std'] = df['close'].rolling(window=20).std()
                df['boll_upper'] = df['boll_mid'] + 2 * df['boll_std']
                df['boll_lower'] = df['boll_mid'] - 2 * df['boll_std']
                
                # 保存数据
                file_path = os.path.join(INDEX_DIR, f'{ts_code}.csv')
                df.to_csv(file_path, index=False)
                print(f"成功下载并保存 {name} 的数据")
            else:
                print(f"未获取到 {name} 的数据")
            
            # 下载指数成分股信息
            if not ts_code.endswith('.BJ'):  # 北证指数暂不支持成分股查询
                try:
                    print(f"下载 {name} 的成分股信息...")
                    components = pro.index_weight(
                        index_code=ts_code,
                        start_date=start_date,
                        end_date=end_date
                    )
                    if not components.empty:
                        components_file = os.path.join(INDEX_DIR, f'{ts_code}_components.csv')
                        components.to_csv(components_file, index=False)
                        print(f"成功保存 {name} 的成分股信息")
                except Exception as e:
                    print(f"下载 {name} 的成分股信息失败: {str(e)}")
            
        except Exception as e:
            print(f"下载 {name} 的数据失败: {str(e)}")
        
        # 避免频繁请求
        time.sleep(1)
    
    print("\n所有指数数据下载完成")

def update_index_data():
    """更新指数数据"""
    if not os.path.exists(INDEX_DIR):
        print("未找到指数数据目录，请先运行完整下载")
        return
    
    print("开始更新指数数据...")
    
    for file in os.listdir(INDEX_DIR):
        if not file.endswith('.csv') or 'components' in file:
            continue
        
        ts_code = file.replace('.csv', '')
        file_path = os.path.join(INDEX_DIR, file)
        
        try:
            # 读取现有数据
            df = pd.read_csv(file_path)
            if not df.empty:
                last_date = df['trade_date'].max()
                start_date = (pd.to_datetime(str(last_date)) + timedelta(days=1)).strftime('%Y%m%d')
                end_date = datetime.now().strftime('%Y%m%d')
                
                print(f"\n更新 {ts_code} 的数据，从 {start_date} 到 {end_date}...")
                
                # 下载新数据
                new_data = pro.index_daily(
                    ts_code=ts_code,
                    start_date=start_date,
                    end_date=end_date
                )
                
                if not new_data.empty:
                    # 合并新旧数据
                    df = pd.concat([df, new_data], ignore_index=True)
                    df = df.drop_duplicates(subset=['trade_date'], keep='last')
                    df = df.sort_values('trade_date', ascending=True)
                    
                    # 重新计算技术指标
                    # 计算均线
                    for ma in [5, 10, 20, 30, 60, 120, 240,250]:
                        df[f'ma{ma}'] = df['close'].rolling(window=ma).mean()
                    
                    # 计算成交量均线
                    for vol_ma in [5, 10, 20, 30, 40, 50, 60, 70, 80, 90, 100, 120, 150, 200, 250]:
                        df[f'vol_ma{vol_ma}'] = df['vol'].rolling(window=vol_ma).mean()
                    
                    # 计算MACD
                    exp1 = df['close'].ewm(span=12, adjust=False).mean()
                    exp2 = df['close'].ewm(span=26, adjust=False).mean()
                    df['macd'] = exp1 - exp2
                    df['macd_signal'] = df['macd'].ewm(span=9, adjust=False).mean()
                    df['macd_hist'] = df['macd'] - df['macd_signal']
                    
                    # 计算RSI
                    delta = df['close'].diff()
                    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
                    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
                    rs = gain / loss
                    df['rsi'] = 100 - (100 / (1 + rs))
                    
                    # 计算布林带
                    df['boll_mid'] = df['close'].rolling(window=20).mean()
                    df['boll_std'] = df['close'].rolling(window=20).std()
                    df['boll_upper'] = df['boll_mid'] + 2 * df['boll_std']
                    df['boll_lower'] = df['boll_mid'] - 2 * df['boll_std']
                    
                    # 保存更新后的数据
                    df.to_csv(file_path, index=False)
                    print(f"成功更新 {ts_code} 的数据")
                else:
                    print(f"{ts_code} 没有新数据需要更新")
        
        except Exception as e:
            print(f"更新 {ts_code} 的数据失败: {str(e)}")
        
        time.sleep(1)
    
    print("\n所有指数数据更新完成")

if __name__ == '__main__':
    # 首次运行，下载所有数据
    ensure_dirs()
    download_index_data()  # 先下载指数数据
    download_all_data()    # 再下载个股数据
    
    # 后续运行，只更新数据
    # update_index_data()
    # update_all_data() 